import pp
import helpers
import os
import pygame
import numpy
import math
import pickle

class map:
	def __init__(self, mapname, level, recalc=False):
		if not recalc:
			try:
				f = file(os.path.join('data', mapname + ".dat"), "r")
				self.gradmap = pickle.load(f)
			except IOError:
				self.gradmap = None
		else:
			self.gradmap = None
		if self.gradmap == None:
			mapimage = helpers.load_image(mapname + ".png",None, True)[0]
			size = (width, height) = (mapimage.get_width(), mapimage.get_height())
			unitsize = (wunits, hunits) = (width / 2**level + 1, height / 2**level + 1)
			chargemap = numpy.empty((width,height),int)
			chargemap = numpy.array(pygame.surfarray.array_alpha(pygame.transform.scale(mapimage, unitsize)))
			self.gradmap = integrate(chargemap, unitsize)
			f = file(os.path.join('data', mapname + ".dat"), "w+")
			pickle.dump(self.gradmap, f)

def intpart(x, chargemap, gradmap, usize):
	line = numpy.zeros((usize[0],2),float)
	for y in range(usize[1]):
		for mx in range(usize[0]):
			for my in range(usize[1]):
				dx = mx-x
				dy = my-y
				if ((dx != 0) and (dy != 0)):
					c = chargemap[mx][my]
					if c != 0:
						r = math.sqrt((dx)**2+(dy)**2)
						line[y][0] += chargemap[mx][my] * 2 * dx / r**3
						line[y][1] += chargemap[mx][my] * 2 * dy / r**3
	if (x % 10) == 0:
		print x
	else:
		print x,
	return line	

def integrate(chargemap, usize):
	clock = pygame.time.Clock()
	clock.tick()
	ppservers = ()
	jobserver = pp.Server(ppservers=ppservers)
	
	lines = []
	gradmap = numpy.zeros((usize[0],usize[1],2),float)
	for x in range(usize[0]):
		lines.append(jobserver.submit(intpart, (x, chargemap, gradmap, usize), (),("numpy","math")))
	
	x = 0
	for line in lines:
		line = line()
		for y in range(usize[1]):
			gradmap[x][y] += line[y]
		x += 1

	print gradmap.max()
	gradmap *= 255.0/gradmap.max()
	clock.tick()
	print "it took ", clock.get_time() / 1000, " s."
	return gradmap
	
	
if __name__ == "__main__":
	pygame.init()
	pygame.display.set_mode((1,1))
	map("map", 3, True)